import pydash

from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver

from app.graph.state import State

from app.graph.NodeChatbot import NodeChatbot
from app.graph.NodeIntentRecog import NodeIntentRecog
from app.graph.NodeWeather import NodeWeather

from app.graph.edge_after_intent_recog import edge_after_intent_recog


def _make_graph():
    g = StateGraph(State)

    g.add_node("intent_recog", NodeIntentRecog())
    g.add_node("task_chatbot", NodeChatbot())
    g.add_node("task_weather", NodeWeather())

    g.add_edge(START, "intent_recog")
    g.add_conditional_edges(
        "intent_recog", edge_after_intent_recog, ["task_chatbot", "task_weather"]
    )
    g.add_edge("task_chatbot", END)
    g.add_edge("task_weather", END)

    return g.compile(checkpointer=False)  # TODO: persist to disk


graph = pydash.memoize(_make_graph)
